test: Add mooncake trace integration test #603
Conversation
Try out this PRQuick install: pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@75ffee6a5a660e95f31d6b1b3971afde35f2982aRecommended with virtual environment (using uv): uv venv --python 3.12 && source .venv/bin/activate
uv pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@75ffee6a5a660e95f31d6b1b3971afde35f2982aLast updated for commit: |
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
debermudez
force-pushed
the
dbermudez/aip-659-add-mooncake-integration-test
branch
from
a628ec6 to
1e5900a
Compare
ajcasagrande
left a comment
ajcasagrande
left a comment
There was a problem hiding this comment.
looks good, just a few notes
WalkthroughA new integration test module for the mooncake_trace custom dataset type is introduced, along with a utility helper function. The test module contains three integration tests validating different Mooncake trace input formats: structured traces with input/output lengths, text-input traces, and multi-turn sessions with session IDs. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@tests/integration/test_mooncake_trace.py`:
- Around line 63-69: The test constructs the list variable traces with entries
containing "text_input" but omits required "hash_ids", causing MooncakeTrace
validation failures; update each dictionary in traces (the list defined as
traces) to include a corresponding "hash_ids" field (e.g., an array or dict
matching the expected MooncakeTrace shape) for every entry that has
"text_input", ensuring the values match the test expectations or use stable
placeholders that satisfy validation.
🧹 Nitpick comments (3)
tests/integration/utils.py (1)
17-36: Avoid blocking file I/O in async test helpers.This helper performs synchronous disk writes; to align with async I/O guidance, consider making it async and offloading the write to a thread (also allows writing bytes in one call). Callers will need to
awaitit. As per coding guidelines, ...♻️ Suggested refactor
+import asyncio import base64 import subprocess from pathlib import Path @@ -def create_mooncake_trace_file( +async def create_mooncake_trace_file( tmp_path: Path, traces: list[dict], filename: str = "traces.jsonl", ) -> Path: @@ """ trace_file = tmp_path / filename - with open(trace_file, "w") as f: - for trace in traces: - f.write(orjson.dumps(trace).decode("utf-8") + "\n") + lines = b"".join(orjson.dumps(trace) + b"\n" for trace in traces) + await asyncio.to_thread(trace_file.write_bytes, lines) return trace_filetests/integration/test_mooncake_trace.py (2)
15-18: Consider parametrizing the three Mooncake trace scenarios.The three test methods share a near-identical CLI invocation; a single parametrized test would reduce duplication and make it easier to add more trace formats. As per coding guidelines, ...
20-25: Add explicit return type hints for async tests.Please annotate async test methods with
-> None(and apply to the other two methods as well). As per coding guidelines, ...✏️ Example update
- async def test_basic_mooncake_trace_with_input_length( + async def test_basic_mooncake_trace_with_input_length( self, cli: AIPerfCLI, aiperf_mock_server: AIPerfMockServer, tmp_path: Path, - ): + ) -> None:
debermudez
deleted the
dbermudez/aip-659-add-mooncake-integration-test
branch
Summary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings.